#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18


import threading
import time
import cv2
import os
import json
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg
from spirecv.base.BaseNode import BaseNode
import argparse
from datetime import datetime


class SpireViewSaverNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown
        )
        self._image_reader = Subscriber(
            '/' + job_name + '/detector/image_results', 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port
        )
        self.output_spire_dir = self.get_param("output_spire_dir", "G:/dataset/240918_false_detections")
        self.save_empty = self.get_param("save_empty", 0)
        self.image_queue = Queue()
        self.result_queue = Queue()
        self.queue_pool.append(self.image_queue)
        self.queue_pool.append(self.result_queue)
        self.run_time = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        self.frame_cnt = 0
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def image_callback(self, msg):
        self.image_queue.put(msg)

    def run(self):
        while self.is_running():
            msg1 = self.image_queue.get(block=True)
            if msg1 is None:
                break

            if 'spirecv_msgs::2DTargets' in msg1:
                self.frame_cnt += 1
                msg2 = msg1['spirecv_msgs::2DTargets']
                image_path = os.path.join(self.output_spire_dir, 'scaled_images')
                anno_path = os.path.join(self.output_spire_dir, 'annotations')
                if not os.path.exists(image_path):
                    os.makedirs(image_path)
                if not os.path.exists(anno_path):
                    os.makedirs(anno_path)

                img = sms2cvimg(msg1)
                spire_anno = {}
                spire_anno['file_name'] = msg2['file_name'] if len(msg2['file_name']) > 0 else "{}_{}.jpg".format(self.run_time, str(self.frame_cnt).zfill(7))
                spire_anno['height'] = img.shape[0]
                spire_anno['width'] = img.shape[1]
                spire_anno['annos'] = []
                for target in msg2['targets']:
                    anno = {}
                    anno['bbox'] = target['bbox']
                    anno['score'] = target['score']
                    anno['category_name'] = target['category_name']
                    anno['area'] = anno['bbox'][2] * anno['bbox'][3]
                    spire_anno['annos'].append(anno)

                if len(spire_anno['annos']) or self.save_empty > 0:
                    output_fn = os.path.join(anno_path, spire_anno['file_name'] + '.json')
                    with open(output_fn, "w") as f:
                        json.dump(spire_anno, f)

                    cv2.imwrite(os.path.join(image_path, spire_anno['file_name']), img)
                    print("Saved: {}".format(spire_anno['file_name']))

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'config',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args = parser.parse_args()
    print("config:", args.config)
    print("job-name:", args.job_name)

    cam = SpireViewSaverNode(args.job_name, param_dict_or_file=args.config)
    cam.join()
